
from langchain_community.chat_models import ChatTongyi
from llama_index.core.indices.query.query_transform.base import HyDEQueryTransform
from llama_index.core.query_engine import TransformQueryEngine
from llama_index.core import Settings,SimpleDirectoryReader,VectorStoreIndex

from llama_index.embeddings.dashscope import (
    DashScopeEmbedding,
    DashScopeTextEmbeddingModels,
    DashScopeTextEmbeddingType
)

#词嵌入模型
embed_model = DashScopeEmbedding(
    model_name=DashScopeTextEmbeddingModels.TEXT_EMBEDDING_V3,
    text_type=DashScopeTextEmbeddingType.TEXT_TYPE_DOCUMENT,
    api_key="sk-f97e3654139742a4b01a99631628d36d"
)

# 初始化LLM
llm = ChatTongyi(model="qwen-plus", api_key="sk-f97e3654139742a4b01a99631628d36d")

Settings.llm = llm
Settings.embed_model = embed_model

hyde = HyDEQueryTransform(include_original=True)

docs = SimpleDirectoryReader("D:\Code\sshcode\RAG_pro\docs").load_data()
index = VectorStoreIndex.from_documents(docs)
query_engine = index.as_query_engine(similarity_top_k=5,streaming=True)
query_engine = TransformQueryEngine(
    query_engine=query_engine,
    query_transform=hyde,
)
question = "张华是什么部门的做什么职务？"
res = query_engine.query(question)
print(res.print_response_stream())
print("-----------------------------")
#查看大模型生成的假设文档
query_bundle = hyde(question)
hyde_doc = query_bundle.embedding_strs[0]
print(hyde_doc)